from PIL import Image
from math import sqrt
import numpy as np
import time
import matplotlib.pyplot as plt


def remap_color(channel: np.ndarray, w: float, offset: float):
    bias = 128 * (1 - w)

    channel = w * channel.astype(np.int32) + bias + offset
    return (np.clip(channel, 0, 255)).astype(np.uint8)


def remap_image_method(img_arr: np.ndarray, w: float, offset: float = 0):
    new_img = np.zeros_like(img_arr, dtype=np.uint8)

    # remap function
    k = w
    for i in range(3):
        new_img[:, :, i] = remap_color(img_arr[:, :, i], k, offset)

    # plt.plot(np.linspace(0, 255, 512), remap_color(np.linspace(0, 255, 512), k))
    # plt.show()

    new_img = Image.fromarray(new_img)
    new_img.show()
    new_img.save("new_img.png")
    time.sleep(3)
    return new_img


def simple_threshold(img_arr: np.ndarray, thresholds: np.ndarray):
    img_arr = (img_arr - img_arr.mean(axis=(0, 1))) / img_arr.std(axis=(0, 1))

    color_means = np.mean(img_arr, axis=(0, 1))
    color_least = np.argmin(color_means)

    # thresholds = np.mean(img_arr, axis=(0, 1))  # + 0.9 * np.min(img_arr, axis=(0, 1))
    thresholds = [0, 0, 0]
    thresholds[color_least] = 0.5

    print(thresholds)

    print(img_arr.shape)
    new_img = np.zeros_like(img_arr, dtype=np.uint8)

    for channel in range(3):
        channel_img = img_arr[:, :, channel]
        pixels = channel_img > thresholds[channel]
        channel_img[pixels] = 255
        channel_img[~pixels] = 0

        new_img[:, :, channel] = channel_img

    new_img = Image.fromarray(new_img)
    new_img.show()
    new_img.save("new_img.png")
    time.sleep(3)
    return new_img


if __name__ == "__main__":
    img_path = "img.jpeg"

    thresholds = np.array([0.2, 0.3, 0.5]) * 255

    img = Image.open(img_path)
    img = img.convert("RGB")
    # img.show()
    # input("Press Enter to continue...")
    img_arr = np.array(img)

    new_img = remap_image_method(img_arr, 8, 0)
    # simple_threshold(img_arr, thresholds)
